/*
This is part of Screenbar, which is based in part on the Litestep
shell source code.

Copyright (C) 2008-2009 Jim Babcock
Copyright (C) 1997-2009 The LiteStep Development Team

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
*/
#include "stdafx.h"

WindowTracker *windowTracker = NULL;

WindowTracker::WindowTracker()
{
	foreground = NULL;
	changed = true;
	
	firstTask = lastTask = NULL;
}

WindowTracker::~WindowTracker()
{
	TaskData *next = NULL;
	for(TaskData *task = firstTask; task; task=next) {
		next = task->next;
		delete task;
	}
	firstTask = lastTask = NULL;
	
	for(map<HWND,WindowData*>::iterator ii=windows.begin(); ii!=windows.end(); ii++)
		delete ii->second;
	
	for(map<HWND,WindowData*>::iterator ii=closedWindows.begin(); ii!=closedWindows.end(); ii++)
		delete ii->second;
}

void WindowTracker::windowCreated(HWND handle)
{
	// Verify that this window both really new and not something we're
	// deliberately ignoring.
	if(windows.find(handle) != windows.end())
		return;
	if(shouldIgnoreWindow(handle))
		return;
	// Create a struct to track this window
	changed = true;
	WindowData *window = new WindowData(handle);
	windows[handle] = window;
	unownedWindows.insert(window);
	
	// Assign this window a task and desk
	window->desk = vwm->deskFromLocation(window->lastScreenPos);
	if(!window->desk)
		window->desk = vwm->rescueOffscreenWindow(window, &window->lastScreenPos);
	
	updateMinimap(window->desk->getIndex());
	
	updateWindowTask(window);
}

void WindowTracker::updateWindow(WindowData *window)
{
	// If this window has been hidden, become a Litestep window, been
	// reparented, etc., stop tracking it
	if(shouldIgnoreWindow(window->handle)) {
		windowDestroyed(window->handle);
		return;
	}
	
	if(window->updateCachedProperties()) {
		changed = true;
		if(window->desk)
			updateMinimap(window->desk->getIndex());
	}
	
	if(!window->minimized)
	{
		VirtualDesktop *newDesk = vwm->deskFromLocation(window->lastScreenPos);
		if(!newDesk)
			newDesk = vwm->rescueOffscreenWindow(window, &window->lastScreenPos);
		
		if(newDesk != window->desk) {
			if(window->desk)
				updateMinimap(window->desk->getIndex());
			updateMinimap(newDesk->getIndex());
		}
		
		updateAndFlag(newDesk, window->desk, changed);
	}
	
	updateWindowTask(window);
}

void WindowTracker::windowDestroyed(HWND handle)
{
	// Check that we were tracking this window
	if(windows.find(handle) == windows.end())
		return;
	changed = true;
	
	WindowData *window = windows[handle];
	
	if(window->desk)
		updateMinimap(window->desk->getIndex());
	
	if(window->task)
		removeWindowFromTask(window);
	
	if(window == foreground)
		foreground = NULL;
	unownedWindows.erase(window);
	
	// After a window closes, we keep track of its last-known state for awhile
	// so taskbar animations of it closing can finish
	windows.erase(handle);
	window->closeTime = getSystemTime();
	closedWindows[handle] = window;
}

void WindowTracker::windowActivated(HWND handle)
{
	// If we get an activation notice for a window we haven't seen before,
	// handle its creation.
	if(windows.find(handle) == windows.end())
		windowCreated(handle);
	
	// Update the focus
	updateFocus();
}

bool WindowTracker::shouldIgnoreWindow(HWND window)
{
	// Skip hidden windows
	if(!IsWindowVisible(window))
		return true;
	
	// Skip anything that has a parent but isn't a popup
	HWND parent = GetParent(window);
	int windowStyle = GetWindowLong(window, GWL_STYLE);
	if(parent && !(windowStyle & WS_POPUP))
		return true;
	
	// Skip windows which are components of the Litestep desktop (including the VWM itself)
	if(GetWindowLongPtr(window, GWLP_USERDATA) == magicDWord)
		return true;

	return false;
}

BOOL CALLBACK WindowTracker::windowChildListBuilder(HWND window, LPARAM param)
{
	vector<HWND> *vec = (vector<HWND>*)param;
	if(!windowTracker->shouldIgnoreWindow(window))
		vec->push_back(window);
	return TRUE;
}

BOOL CALLBACK WindowTracker::windowListBuilder(HWND window, LPARAM param)
{
	vector<HWND> *vec = (vector<HWND>*)param;
	if(!windowTracker->shouldIgnoreWindow(window))
		vec->push_back(window);
	
	// If this is a root window, iterate over its children
	HWND parent = GetParent(window);
	if(!parent)
		EnumChildWindows(window, &windowChildListBuilder, param);
	
	return TRUE;
}

bool WindowTracker::update()
{
	changed = false;
	
	// Enumerate windows
	zOrder.clear();
	EnumWindows(&windowListBuilder, (LPARAM)&zOrder);
	
	// Check for created and destroyed windows
	vector<HWND> newWindows;
	vector<HWND> deletedWindows;
	
	for(map<HWND,WindowData*>::iterator ii=windows.begin(); ii!=windows.end(); ii++)
		ii->second->visited = false;
	
	for(unsigned ii=0; ii<zOrder.size(); ii++) {
		HWND handle = zOrder[ii];
		if(windows.find(handle)==windows.end())
			newWindows.push_back(handle);
		else
			windows[handle]->visited = true;
	}
	
	for(map<HWND,WindowData*>::iterator ii=windows.begin(); ii!=windows.end(); ii++) {
		if(!ii->second->visited)
			deletedWindows.push_back(ii->second->handle);
	}
	
	// Handle deleted windows
	for(unsigned ii=0; ii<deletedWindows.size(); ii++)
		windowDestroyed(deletedWindows[ii]);
	
	// Check for changes to known windows
	// This may lead to windows being destroyed (if we notice that a window
	// being updated isn't around anymore), so iterate through a copy of the
	// list rather than the original.
	vector<WindowData*> windowsToUpdate(windows.size());
	int pos = 0;
	for(map<HWND,WindowData*>::iterator ii=windows.begin(); ii!=windows.end(); ii++)
		windowsToUpdate[pos++] = ii->second;
	for(unsigned ii=0; ii<windowsToUpdate.size(); ii++)
		updateWindow(windowsToUpdate[ii]);
	
	// Handle newly created windows
	for(unsigned ii=0; ii<newWindows.size(); ii++)
		windowCreated(newWindows[ii]);
	
	// Update focus
	updateFocus();
	
	// Check memory of closed windows for staleness
	// (We remember closed windows so animations can finish)
	UINT64 systemTime = getSystemTime();
	for(map<HWND,WindowData*>::iterator ii = closedWindows.begin(); ii != closedWindows.end(); )
	{
		WindowData *window = ii->second;
		if(window->closeTime + settings->closedWindowMemoryDuration*10000 < systemTime) {
			delete ii->second;
			ii = closedWindows.erase(ii);
		} else {
			ii++;
		}
	}
	
	return changed;
}

void WindowTracker::updateFocus()
{
	HWND pos = GetForegroundWindow();
	
	if(windows.find(pos) == windows.end())
		windowCreated(pos);
	
	if(windows.find(pos) != windows.end()) {
		WindowData *window = windows[pos];
		if(foreground != window) {
			changed = true;
			if(window->desk)
				updateMinimap(window->desk->getIndex());
		}
		foreground = windows[pos];
	} else {
		if(foreground != NULL) {
			changed = true;
			if(foreground->desk)
				updateMinimap(foreground->desk->getIndex());
		}
		foreground = NULL;
	}
}

bool WindowTracker::checkForChange()
{
	if(changed) {
		changed = false;
		return true;
	} else {
		return false;
	}
}

void WindowTracker::updateMinimap(int desk)
{
	while(desk >= (int)minimapChangeIndices.size())
		minimapChangeIndices.push_back(0);
	minimapChangeIndices[desk]++;
}

int WindowTracker::getMinimapChangeIndex(int desk)
{
	if(desk < (int)minimapChangeIndices.size())
		return minimapChangeIndices[desk];
	else
		return 0;
}

TaskData *WindowTracker::getFirstTask()
{
	return firstTask;
}

TaskData *WindowTracker::getLastTask()
{
	return lastTask;
}

WindowData *WindowTracker::getForegroundWindow()
{
	return foreground;
}

TaskData *WindowTracker::getForegroundTask()
{
	WindowData *foregroundWindow = getForegroundWindow();
	if(!foregroundWindow)
		return NULL;
	return foregroundWindow->task;
}

/// If the focused window is on another desktop or is the desktop itself, focus
/// on the topmost window of this desktop. If there is no such window, focus on
/// the desktop.
void WindowTracker::raiseLocalForeground()
{
	for(unsigned ii=0; ii<zOrder.size(); ii++)
	{
		HWND handle = zOrder[ii];
		if(windows.find(handle) == windows.end()) continue;
		WindowData *window = windows[handle];
		
		// Only focus something that's on a visible desk
		if(!window->desk->isVisible())
			continue;
		// Don't focus minimized windows (that'd be silly)
		if(window->isMinimized())
			continue;
		
		vwm->raiseWindow(window);
		break;
	}
	
	// If we got throught he whole window list without finding something
	// suitable to focus, focus on the desktop.
	SetFocus(getDesktopWindow());
}

void WindowTracker::getTasks(VirtualDesktop *desk, vector<TaskData*> &outTasks)
{
	if(desk)
	{
		for(TaskData *task=firstTask; task!=NULL; task=task->next)
		{
			vector<VirtualDesktop*> desks;
			task->getDesks(desks);
			for(unsigned jj=0; jj<desks.size(); jj++) {
				if(desks[jj] == desk) {
					outTasks.push_back(task);
					break;
				}
			}
		}
	}
	else
	{
		for(TaskData *task=firstTask; task!=NULL; task=task->next)
			outTasks.push_back(task);
	}
}

void WindowTracker::getWindows(VirtualDesktop *desk, vector<WindowData*> &outWindows)
{
	for(unsigned ii=0; ii<zOrder.size(); ii++)
	{
		HWND handle = zOrder[ii];
		if(windows.find(handle) == windows.end())
			continue;
		WindowData *window = windows[handle];
		if(!desk || window->desk == desk)
			outWindows.push_back(window);
	}
}

WindowData *WindowTracker::getWindow(HWND window)
{
	map<HWND,WindowData*>::iterator it = windows.find(window);
	if(it != windows.end())
		return it->second;
	
	it = closedWindows.find(window);
	if(it != closedWindows.end())
		return it->second;
	
	return NULL;
}

bool WindowTracker::deskIsEmpty(VirtualDesktop *desk)
{
	for(unsigned ii=0; ii<zOrder.size(); ii++)
	{
		HWND handle = zOrder[ii];
		if(windows.find(handle) == windows.end())
			continue;
		WindowData *window = windows[handle];
		if(window->desk == desk)
			return false;
	}
	
	return true;
}

void WindowTracker::reorderTask(TaskData *task, TaskData *prev, TaskData *next)
{
	assert(task);
	assert(prev || next);
	
	if(prev && task->prev == prev)
		return;
	if(next && task->next == next)
		return;
	
	// Unlink the task from the list
	if(task->prev) task->prev->next = task->next;
	if(task->next) task->next->prev = task->prev;
	if(firstTask == task) firstTask = task->next;
	if(lastTask == task) lastTask = task->prev;
	
	if(prev) {
		// The preceding task was given, so insert into the list after it
		task->prev = prev;
		task->next = prev->next;
		if(task->next)
			task->next->prev = task;
		prev->next = task;
		if(lastTask == prev)
			lastTask = task;
	} else {
		// The following task was given, so insert into the list before it
		task->next = next;
		task->prev = next->prev;
		if(task->prev)
			task->prev->next = task;
		next->prev = task;
		if(firstTask == next)
			firstTask = task;
	}
}

void WindowTracker::updateWindowTask(WindowData *window)
{
	if(window->task)
	{
		TaskData *task = window->task;
		WindowData *representative = task->getRepresentativeWindow();
		
		// If this window is the representative of a task, verify that it is
		// still eligible for that role
		if(window == representative)
		{
			if(!window->isTask)
				removeTask(task);
		}
		// If this window is the child of a task, check that this is still
		// correct.
		else if(window->isTask || findParentTaskWindow(window) != representative)
		{
			removeWindowFromTask(window);
			updateWindowTask(window);
		}
	}
	else
	{
		// This window isn't currently considered part of any task, and is
		// eligible to be a task by itself, so make it one.
		if(window->isTask)
		{
			// Create a task and give it ownership of this window
			TaskData *task = new TaskData(window);
			window->task = task;
			unownedWindows.erase(window);
			
			// Add this task to our data structures
			if(lastTask) {
				lastTask->next = task;
				task->prev = lastTask;
				task->next = NULL;
				lastTask = task;
			} else {
				firstTask = lastTask = task;
				task->prev = task->next = NULL;
			}
			tasksByPid.insert(pair<DWORD,TaskData*>(window->pid, task));
			
			// Look through windows that aren't currently owned by a task, and
			// claim any that belong to this window.
			for(set<WindowData*>::iterator ii=unownedWindows.begin(); ii!=unownedWindows.end(); ii++)
			{
				WindowData *child = *ii;
				if(findParentTaskWindow(child) == window) {
					child->task = task;
					task->addWindow(child);
					unownedWindows.erase(child);
					break;
				}
			}
		}
		// This window isn't currently part of any task and isn't eligible to
		// be a task by itself, so look for tasks it could be considered part
		// of.
		else
		{
			WindowData *parentTaskWindow = findParentTaskWindow(window);
			if(parentTaskWindow && parentTaskWindow->task)
			{
				TaskData *task = parentTaskWindow->task;
				
				window->task = task;
				task->addWindow(window);
				unownedWindows.erase(window);
			}
		}
	}
}

void WindowTracker::removeTask(TaskData *task)
{
	vector<WindowData*> windows;
	task->getWindows(windows);
	WindowData *representative = task->representative;
	
	// Mark the windows that were owned by this task as taskless
	for(unsigned ii=0; ii<windows.size(); ii++)
		windows[ii]->task = NULL;
	
	// Unlink this task
	if(firstTask==task) firstTask = task->next;
	if(lastTask==task) lastTask = task->prev;
	if(task->prev) task->prev->next = task->next;
	if(task->next) task->next->prev = task->prev;
	
	// Remove this task from the PID->task mapping
	DWORD pid = task->representative->pid;
	for(multimap<DWORD,TaskData*>::iterator ii = tasksByPid.find(pid); ii!=tasksByPid.end(); ii++)
	{
		if(ii->second == task) {
			tasksByPid.erase(ii);
			break;
		}
		if(ii->first != pid)
			break;
	}
	
	// Delete the task
	delete task;
	
	// Find new tasks to assign ownership of windows to
	for(unsigned ii=0; ii<windows.size(); ii++) {
		if(windows[ii] != representative)
			updateWindowTask(windows[ii]);
	}
}

void WindowTracker::removeWindowFromTask(WindowData *window)
{
	TaskData *task = window->task;
	if(!task) return;
	
	window->task = NULL;
	
	if(window == task->representative)
		removeTask(task);
	else
		task->removeWindow(window);
}

WindowData *WindowTracker::findParentTaskWindow(WindowData *window)
{
	// Non-task windows are assigned to the most-closely-related task found.
	// The main effect of this is that these windows move between desks with
	// the window they're assigned to. The preference order is:
	//   - The window's parent
	//   - The window's owner
	//   - A task window owned by the same PID and on the same monitor
	//   - A task window owned by the same PID but not on the same monitor
	// Associating non-task windows with task windows owned by the same
	// process helps keep related windows together when moving tasks between
	// desktops. In many cases, this is the only indication that the windows
	// are related; eg, Photoshop tool windows and Winamp playlist/equalizer
	// windows. Favoring task windows on the same desk prevents an issue
	// where you have a single process with task windows on two monitors, and
	// one of them makes a tooltip, causing the other to be temporarily
	// mismarked as spanning both monitors.
	HWND parentHandle = GetParent(window->handle);
	HWND ownerHandle = GetWindow(window->handle, GW_OWNER);
	if(windows.find(parentHandle) != windows.end()) {
		WindowData *parent = windows[parentHandle];
		if(parent->isTask) return parent;
	}
	if(windows.find(ownerHandle) != windows.end()) {
		WindowData *owner = windows[ownerHandle];
		if(owner->isTask) return owner;
	}
	
	vector<TaskData*> pidCandidates;
	for(multimap<DWORD,TaskData*>::iterator iter = tasksByPid.find(window->pid);
	    iter != tasksByPid.end() && iter->first==window->pid; iter++)
	{
		pidCandidates.push_back(iter->second);
	}
	
	for(size_t ii=0; ii<pidCandidates.size(); ii++)
	{
		if(pidCandidates[ii]->getRepresentativeWindow()->desk == window->desk)
			return pidCandidates[ii]->getRepresentativeWindow();
	}
	if(pidCandidates.size())
		return pidCandidates[0]->getRepresentativeWindow();
	
	return NULL;
}

/////////////////////////////////////////////////////////////////////////////

TaskData::TaskData(WindowData *representative)
{
	this->representative = representative;
}

void TaskData::getWindows(vector<WindowData*> &outWindows)
{
	outWindows.push_back(representative);
	outWindows.insert(outWindows.end(), children.begin(), children.end());
}

void TaskData::getWindows(VirtualDesktop *desk, vector<WindowData*> &outWindows)
{
	if(representative->desk == desk)
		outWindows.push_back(representative);
	for(unsigned ii=0; ii<children.size(); ii++) {
		if(children[ii]->desk == desk)
			outWindows.push_back(children[ii]);
	}
}

void TaskData::getDesks(vector<VirtualDesktop*> &outDesks)
{
	// Clear the visited flag of desks this task is on
	for(unsigned ii=0; ii<children.size(); ii++)
		children[ii]->desk->visited = false;
	
	// Add the representative window's desk
	representative->desk->visited = true;
	outDesks.push_back(representative->desk);
	
	// Iterate through windows, adding desks and setting the visited flag so
	// that we don't add the same desk twice
	for(unsigned ii=0; ii<children.size(); ii++)
	{
		if(!children[ii]->desk->visited) {
			outDesks.push_back(children[ii]->desk);
			children[ii]->desk->visited = true;
		}
	}
}

WindowData *TaskData::getRepresentativeWindow()
{
	return representative;
}

HICON TaskData::getIcon(VirtualDesktop *desk, int size)
{
	return representative->getIcon(size);
}

bool TaskData::containsWindow(WindowData *window)
{
	if(window==representative)
		return true;
	
	for(unsigned ii=0; ii<children.size(); ii++) {
		if(children[ii] == window)
			return true;
	}
	return false;
}

void TaskData::addWindow(WindowData *window)
{
	if(!containsWindow(window))
		children.push_back(window);
}

void TaskData::removeWindow(WindowData *window)
{
	for(vector<WindowData*>::iterator ii=children.begin(); ii!=children.end(); ii++)
	{
		if(*ii == window) {
			children.erase(ii);
			break;
		}
	}
}

